"""
将 南方医 数据集转换成 yolo 训练集
http://ludo17.free.fr/mitos_2012/index.html

"""
import json
import os

import numpy as np
import openslide
from PIL import Image
import cv2
from tqdm import tqdm

if __name__ == '__main__':
    # GIST已审核-127中   标注了总共1528个有丝分裂细胞
    wsi_path = '/media/hsmy/wanghao_18T/有丝分裂/归档文件/南方医/源数据/GIST已审核_WSI/'
    geo_path = '/media/hsmy/wanghao_18T/有丝分裂/归档文件/南方医/源数据/GIST已审核-127/'

    png_path = '/media/hsmy/wanghao_18T/dataset/南方医/yolo_40_1024_origin_/images/train/'
    label_path = '/media/hsmy/wanghao_18T/dataset/南方医/yolo_40_1024_origin_/labels/train/'

    mask_path = '/media/hsmy/wanghao_18T/dataset/南方医/yolo_40/masks/'
    os.makedirs(mask_path, exist_ok=True)

    os.makedirs(png_path, exist_ok=True)
    os.makedirs(label_path, exist_ok=True)

    patch_size = 1024  # 训练 image 大小
    patch_level = 0  # 0：40X 1：10X
    react_size = 64  # 框的大小 40X 下是64x64

    if patch_level != 0:
        react_size = react_size // (patch_level * 4)

    react_size_percent = react_size / patch_size
    for file_path in tqdm(os.listdir(wsi_path)):
        file_name = file_path.split('.')[0]

        slide = openslide.open_slide(os.path.join(wsi_path, file_path))
        width, height = slide.level_dimensions[patch_level]

        # *2变为20X
        # width = width * 2
        # height = height * 2

        mask = np.zeros((height, width), dtype=np.uint8)

        json_path = os.path.join(geo_path, file_name + '.geojson')
        json_obj = json.load(open(json_path))
        features_arr = json_obj['features']

        for index, feature in enumerate(features_arr):
            arr = feature['geometry']['coordinates'][0]
            arr = np.array(arr)
            x = int(np.mean(arr[:, 0]))
            y = int(np.mean(arr[:, 1]))
            if patch_level != 0:
                x = x // (patch_level * 4)
                y = y // (patch_level * 4)
            mask[y, x] = 1

        index = 0
        for x in range(0, width, patch_size):
            if x + patch_size > width:
                x = width - patch_size

            for y in range(0, height, patch_size):
                if y + patch_size > height:
                    y = height - patch_size  # 保证patch规格，超出边界往前推

                # 定义块的范围
                y_end = min(y + patch_size, height)
                x_end = min(x + patch_size, width)

                mask_patch = mask[y:y_end, x:x_end]

                if mask_patch.max() == 0:
                    continue

                # 切割块 location: (x, y) tuple giving the top left pixel in the level 0
                if patch_level != 0:
                    png_patch = slide.read_region((x * 4 * patch_level, y * 4 * patch_level), patch_level,
                                                  (patch_size, patch_size)).convert('RGB')
                else:
                    png_patch = slide.read_region((x, y), patch_level, (patch_size, patch_size)).convert('RGB')

                # png_patch = png_patch.resize((patch_size, patch_size), Image.LANCZOS)  # 由10X扩大为20X

                png_name = f"{file_name}_{index}.png"
                txt_name = f"{file_name}_{index}.txt"

                points = np.argwhere(mask_patch == 1)
                for point in points:
                    label = (
                        0,
                        point[1] / patch_size,
                        point[0] / patch_size,
                        react_size_percent,
                        react_size_percent
                    )
                    with open(label_path + txt_name, 'a') as f:
                        f.write(('%g ' * len(label)).rstrip() % label + '\n')

                png_patch.save(png_path + png_name)

                # kernel = np.ones((10, 10), np.uint8)  # 定义膨胀结构元素
                # diffused_matrix = cv2.dilate(mask_patch, kernel, iterations=1)
                # Image.fromarray((diffused_matrix * 255).astype(np.uint8)).save(mask_path + png_name)
                index += 1
